///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2015-2016, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///

#ifndef S2E_PLUGINS_GENERATIONSEARCHER_H
#define S2E_PLUGINS_GENERATIONSEARCHER_H

#include <s2e/CorePlugin.h>
#include <s2e/Plugin.h>
#include <s2e/Plugins/Core/BaseInstructions.h>
#include <s2e/Plugins/OSMonitors/Support/ModuleExecutionDetector.h>
#include <s2e/Plugins/ExecutionTracers/TestCaseGenerator.h>
#include <s2e/Plugins/Searchers/MultiSearcher.h>
#include <s2e/Plugins/StaticAnalysis/ControlFlowGraph.h>
#include <s2e/Plugins/Support/KeyValueStore.h>
#include <s2e/S2EExecutionState.h>
#include <s2e/Synchronization.h>
#include <s2e/Utils.h>
#include <klee/util/ExprUtil.h>
#include <klee/Searcher.h>

#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/DenseMap.h>
#include <memory>
#include <queue>
#include <deque>
#include <random>
#include <string>
#include "CUPASearcher.h"

namespace s2e {
namespace plugins {
namespace generations {

struct Range {
	uint64_t start;
	uint64_t end;
	bool operator()(const Range &p1, const Range &p2) const {
		return p1.start < p2.start;
	}
	bool operator==(const Range &p1) const {
		return p1.start == start && p1.end == end;
	}
};
typedef std::set<Range, Range> RangeEntries;

typedef std::vector<klee::ref<klee::Expr>> ExprList;

typedef std::map<uint64_t /* hitpoint */, ExprList /* Constraints */> MultiHitConstraints;

typedef std::map<uint64_t /* hitpoint */, int /* count */> MultiHitRecord;

class GenerationSearcher : public Plugin, public klee::Searcher{
    S2E_PLUGIN

public:

    GenerationSearcher(S2E *s2e) : Plugin(s2e) {
    }
    ~GenerationSearcher();
    void initialize();

    virtual klee::ExecutionState &selectState();
    virtual void update(klee::ExecutionState *current, const klee::StateSet &addedStates,
                        const klee::StateSet &removedStates);

    virtual bool empty();

private:
    typedef std::deque<uint64_t> ForkHistory;
    typedef llvm::DenseSet<uint64_t> UnForkInst;
    typedef llvm::DenseSet<uint64_t> OnlyForkConditions;
    typedef llvm::DenseSet<uint64_t> DisableForkConditions;
    typedef llvm::DenseSet<uint64_t> OnlyForkPoints;
    typedef std::set<S2EExecutionState *> States;
    typedef std::pair<std::string, std::vector<unsigned char>> VarValuePair;
    typedef std::vector<VarValuePair> ConcreteInputs;
    typedef llvm::DenseMap<uint64_t, uint64_t> ForkCounts;
    typedef std::map<std::string, ForkCounts> ModuleForkCounts;

    MultiSearcher *m_multiSearcher;
    CUPASearcher *m_cupa;
    ModuleExecutionDetector *m_detector;
    testcases::TestCaseGenerator *m_tc;
    /// Stores the pointer to state 0
    S2EExecutionState *m_initialState;

    MultiHitConstraints multiHitConstraints;
    MultiHitRecord multiHitRecord;

    ModuleForkCounts m_forkCount;
    uint64_t m_maxUnForkGen;

    bool m_underConstraint;
    bool m_searchBug;
    bool m_genMultiHitCase;

    ForkHistory m_forkHistory;
    ForkHistory m_currentForkReocord;
    UnForkInst m_unForkInstCache;

    OnlyForkConditions m_onlyForkConditions;
    DisableForkConditions m_disableForkConditions;
	OnlyForkPoints m_onlyForkPoints;
	OnlyForkPoints m_configedOnlyForkPoints;

    RangeEntries m_forkRanges;//cherry
    void onSymbolicAddress(S2EExecutionState *state, klee::ref<klee::Expr> virtualAddress, uint64_t concreteAddress,
                           bool &concretize, CorePlugin::symbolicAddressReason reason);
    void onStateForkDecide(S2EExecutionState *state, bool *doFork);
    void onUnableFork(S2EExecutionState *state, const klee::ref<klee::Expr> &condition);
    void switchToCUPA();
    void switchToGenerationSearcher();
    void onStateFork(S2EExecutionState *oldState, const std::vector<S2EExecutionState *> &newStates,
                     const std::vector<klee::ref<klee::Expr>> &);
    void onStateKill(S2EExecutionState *state);

    std::string generateJsonHistoryFile(S2EExecutionState *state);
    std::string generateNotForkJsonHistoryFile(S2EExecutionState *state, const std::string &name_id);
    void generateOffsetInfo(ExprList elist, const std::string &name_id);
    void generateJsonHistoryFile(S2EExecutionState *state, const std::string &filePath);
    void generateJsonHistory(S2EExecutionState *state, std::stringstream &ss);

    void generationBugCase(S2EExecutionState *state, klee::ref<klee::Expr> symbolicAddress, uint64_t currentAddress);
    bool solveConstraints(S2EExecutionState *state, const klee::ref<klee::Expr> condition_e, klee::Assignment &assignment);
    void onModuleTranslateBlockEnd(ExecutionSignal *signal, S2EExecutionState *state, const ModuleDescriptor &module,
                                       TranslationBlock *tb, uint64_t endPc, bool staticTarget, uint64_t targetPc);
    void onModuleTranslateBlockStart(ExecutionSignal *signal, S2EExecutionState *state, const ModuleDescriptor &module,
                                     TranslationBlock *tb, uint64_t pc);
    void onOnlyForkTargetExecution(S2EExecutionState *state, uint64_t pc);
    void onDisableForkTargetExecution(S2EExecutionState *state, uint64_t pc);
public:

    ///
    /// \brief isAvailable returns availability of the searcher on the current S2E instance
    ///
    /// Only the S2E instance that has the state 0 has the generation searcher available. Other
    /// instances cannot fetch generations because they don't have state 0.
    ///
    /// \return the availability status
    ///
    bool isAvailable() const {
        return m_initialState != nullptr;
    }

	void generateTestCaseForExprList(S2EExecutionState *state, ExprList explist, const std::string &prefix);

};

} // namespace generations
} // namespace plugins
} // namespace s2e

#endif // S2E_PLUGINS_GENERATIONSEARCHER_H
